# -*-coding:utf-8-*-
import json
import pandas as pd
import numpy as np
import redis
import time
import requests
import sys
from datetime import datetime
from websocket import create_connection

pool = redis.ConnectionPool(host='localhost',
                            port=6379,
                            decode_responses=True,
                            db=0)
r = redis.Redis(connection_pool=pool)

symbol = sys.argv[1]  # ethusdt
freq = sys.argv[2]  # 1000ms
n = sys.argv[3]  # 10
n = int(n)
tp = sys.argv[4]  # 'redis' or 'print'
save_redis = False
if_print = False
if tp == 'redis':
    save_redis = True
if tp == 'print':
    if_print = True

url = f'https://api.binance.com/api/v3/depth?symbol={symbol.upper()}&limit=1000'
res = requests.get(url)
data = json.loads(res.text.encode())
snapshot_updateId = data['lastUpdateId']
bids = np.array(data['bids']).astype(float)
asks = np.array(data['asks']).astype(float)
bids = dict(zip(bids[:, 0], bids[:, 1]))
asks = dict(zip(asks[:, 0], asks[:, 1]))

sorted_ask_price = sorted(asks.keys())[:n]
ask_vols = [asks.get(x) for x in sorted_ask_price]
sorted_bid_price = sorted(bids.keys())[:n]
bid_vols = [asks.get(x) for x in sorted_bid_price]

ws = create_connection(f"wss://stream.binance.com:9443/ws/{symbol.lower()}@depth@{freq}")

last_sorted_ask_price = np.array([0] * n)
last_sorted_bid_price = np.array([0] * n)
last_ask_vols = np.array([0] * n)
last_bid_vols = np.array([0] * n)

count = 0
while True:
    count += 1
    result = ws.recv()
    result = json.loads(result)
    ask = np.array(result['a']).astype(float)
    bid = np.array(result['b']).astype(float)
    E = result['E']
    ts = E / 1000
    t = datetime.fromtimestamp(ts)
    str_t = datetime.strftime(t, '%Y-%m-%d %H:%M:%S')

    if len(bid) > 0:
        bid = dict(zip(bid[:, 0], bid[:, 1]))
        bids.update(bid)

        bids0 = np.array([x for x in bids.items() if x[1] == 0])
        bids = np.array([x for x in bids.items() if x[1] != 0])

        bids = dict(zip(bids[:, 0], bids[:, 1]))
        sorted_bid_price = sorted(bids.keys(), reverse=True)[:n]
        bid_vols = [bids.get(x) for x in sorted_bid_price]

    if len(ask) > 0:
        ask = dict(zip(ask[:, 0], ask[:, 1]))
        asks.update(ask)
        asks = np.array([x for x in asks.items() if x[1] != 0])
        asks = dict(zip(asks[:, 0], asks[:, 1]))
        sorted_ask_price = sorted(asks.keys())[:n]
        ask_vols = [asks.get(x) for x in sorted_ask_price]
        sorted_ask_price = sorted_ask_price[::-1]
        ask_vols = ask_vols[::-1]

    # asks 是降序的, bids 是升序的
    sorted_ask_price = np.array(sorted_ask_price)
    sorted_bid_price = np.array(sorted_bid_price)
    ask_vols = np.array(ask_vols)
    bid_vols = np.array(bid_vols)

    if if_print:
        print('---')
        for i in range(n):
            print('%.02f' % sorted_ask_price[i], '%.05f' % ask_vols[i])

        print('---')

        for i in range(n):
            print('%.02f' % sorted_bid_price[i], '%.05f' % bid_vols[i])

        print('\n\n')

    if save_redis:
        depth_data = {'t': str_t,
                      'ask_price': list(sorted_ask_price), 'ask_vols': list(ask_vols),
                      'bid_price': list(sorted_bid_price), 'bid_vols': list(bid_vols)}
        print(depth_data)
        depth_data = json.dumps(depth_data)
        r.set(f'{symbol}_depth', depth_data)

    last_sorted_ask_price = sorted_ask_price
    last_sorted_bid_price = sorted_bid_price
    last_ask_vols = ask_vols
    last_bid_vols = bid_vols
